package storage

import (
	"context"

	"gitcode.com/eteam/sae.git/types"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

type MessageInfoDbItem struct {
	MessageId string            `bson:"_id"`
	TaskId    string            `bson:"taskId"`
	TimeStamp int64             `bson:"timeStamp"`
	Role      types.MessageRole `bson:"role"`
	Parts     []types.Part      `bson:"parts"`
}

func (item *MessageInfoDbItem) ToMessage() *types.Message {
	if item == nil {
		return nil
	}
	return &types.Message{
		MessageId: item.MessageId,
		TimeStamp: item.TimeStamp,
		Role:      item.Role,
		Parts:     item.Parts,
		Kind:      types.MessageObjectKind,
	}
}

type _MessageInfoDao struct {
	mc *mongo.Collection
}

func newMessageInfoDao(db *mongo.Database) (*_MessageInfoDao, error) {
	mc := db.Collection("message_info")
	_, err := mc.Indexes().CreateOne(context.Background(), mongo.IndexModel{
		Keys: bson.D{{"taskId", 1}, {"timeStamp", 1}},
	})
	if err != nil {
		return nil, err
	}
	return &_MessageInfoDao{mc: mc}, nil
}

func (dao *_MessageInfoDao) Insert(ctx context.Context, dbItem *MessageInfoDbItem) error {
	_, err := dao.mc.InsertOne(ctx, dbItem)
	return err
}

func (dao *_MessageInfoDao) Upsert(ctx context.Context, dbItem *MessageInfoDbItem) error {
	_, err := dao.mc.UpdateOne(ctx, bson.M{"taskId": dbItem.TaskId, "_id": dbItem.MessageId},
		bson.M{"$set": dbItem}, options.UpdateOne().SetUpsert(true))
	return err
}

func (dao *_MessageInfoDao) ListByTask(ctx context.Context, taskIdList []string) ([]*MessageInfoDbItem, error) {
	if len(taskIdList) == 0 {
		return []*MessageInfoDbItem{}, nil
	}
	cursor, err := dao.mc.Find(ctx, bson.M{"taskId": bson.M{"$in": taskIdList}},
		options.Find().SetSort(bson.M{"timeStamp": 1}))
	if err != nil {
		return nil, err
	}
	retList := []*MessageInfoDbItem{}
	err = cursor.All(ctx, &retList)
	return retList, err
}
